/*
 * Project CelestialBeing: unique_object.hpp
 * Created: 2023/12/15
 * This program uses MulanPSL2. See LICENSE for more.
 */

// This file defines the base class for all unique objects.
// Unique objects are objects that have a unique id. They are not copyable nor movable.
// They are registered in a global registry, and can be accessed by their id.
// They can be pointed to by serializable pointers.

#ifndef CELESTIALBEING_UNIQUE_OBJECT_HPP
#define CELESTIALBEING_UNIQUE_OBJECT_HPP
// gnu g++ needs this
#include <mutex>
#include <shared_mutex>
#include <atomic>
#include "cb/common/celestial_being_base.hpp"
#include "singleton.hpp"
#include "base_classes.hpp"
#include "serializable.hpp"
CB_NAMESPACE_BEGIN

// I want to do things like "class UniqueObjectID: uint64_t", however, base class must be a "class / struct".
// Resort to this wrapper. --- Hineven 2023/11/28

namespace detail {
    // Don't use it, use UniqueID instead.
    class UniqueObjectID_Wrapper {
    public:
        enum Value : uint64_t {
            kInvalid = UINT64_MAX
        };
    };
}

// ID type used among states.
typedef detail::UniqueObjectID_Wrapper::Value UniqueObjectID;

// Thus we can use UniqueObjectID::kInvalid to represent an invalid id.
// "UniqueObjectID id;" is valid also. Explicit conversion is required for uint64_t and UniqueObjectID cast.

class UniqueObject;

template<typename T> concept CUniqueObject = std::is_base_of<UniqueObject, T>::value;

// The global registry for unique objects.
// It is a singleton. It is thread-safe.
class UniqueObjectRegistry : public Singleton<UniqueObjectRegistry>, public NonCopyable, public NonMovable {
    typedef std::shared_lock<std::shared_mutex> ReaderLock;
    typedef std::unique_lock<std::shared_mutex> WriterLock;
public:

    template<CUniqueObject T, CUint64ID<UniqueObjectID> IDType>
    std::shared_ptr<T> GetObject (IDType id) const {
        ReaderLock lock(object_id_map_mutex_);
        auto iter = object_id_map_.find(reinterpret_cast<UniqueObjectID>(id));
        if(iter == object_id_map_.end() || iter->second.expired()) return nullptr;
        // Don't remove the iter if expired here. It may cause deadlocks.
        auto ptr = iter->second.lock();
        auto return_ptr = std::dynamic_pointer_cast<T>(ptr);
        // Casting failure denotes a type dis-match.
        assert((!ptr || return_ptr) && "Potentially an invalid typecast, check the types of the objects.");
        return return_ptr;
    }

    template<CUniqueObject T, typename...Args>
    std::shared_ptr<T> CreateAndRegisterObjectWithID (UniqueObjectID id, Args...args) {
        auto object = std::shared_ptr<T>(new T(args...));
        object->id_ = id;
        RegisterObject(object, id);
        return object;
    }

    template<CUniqueObject T, typename...Args>
    std::shared_ptr<T> CreateAndRegisterObject (Args...args) {
        auto object = std::shared_ptr<T>(new T(args...));
        object->id_ = AllocateID();
        RegisterObject(object, object->id_);
        return object;
    }

    inline bool Exists (UniqueObjectID id) const {
        ReaderLock lock(object_id_map_mutex_);
        return object_id_map_.find(id) != object_id_map_.end();
    }

protected:
    UniqueObjectID AllocateID () {
        return static_cast<UniqueObjectID>(next_id_.fetch_add(1));
    }

    template<CUniqueObject T>
    UniqueObjectID RegisterObject (std::shared_ptr<T> object, UniqueObjectID id = UniqueObjectID::kInvalid) {
        if(id == UniqueObjectID::kInvalid) id = AllocateID();
        assert(id != UniqueObjectID::kInvalid);
        WriterLock lock(object_id_map_mutex_);
        assert(object_id_map_.insert(std::make_pair(id, std::static_pointer_cast<UniqueObject>(object))).second);
        return id;
    }
    // Dangerous function. We don't allow the user to unregister any object manually.
//    void UnregisterObject (UniqueObjectID id) {
//        WriterLock lock(object_id_map_mutex_);
//        auto it = object_id_map_.find(id);
//        assert(it != object_id_map_.end());
//        object_id_map_.erase(it);
//    }

private:
    UniqueObjectRegistry() {}
    ~UniqueObjectRegistry() {}
    friend class Singleton<UniqueObjectRegistry>;

    std::atomic<uint64_t > next_id_;
    std::unordered_map<UniqueObjectID, std::weak_ptr<UniqueObject>> object_id_map_;
    mutable std::shared_mutex object_id_map_mutex_;
};

// Use UniqueObjectRegistry::CreateAndRegisterObject to create a unique object.
class UniqueObject {
public:
    UniqueObjectID GetObjectID () const {
        return id_;
    }
    friend class UniqueObjectRegistry;
protected:
    UniqueObject () {}
    UniqueObject (UniqueObjectID id) : id_(id) {
        assert(id != UniqueObjectID::kInvalid);
    }
    virtual ~UniqueObject() {}
private:
    UniqueObjectID id_;
};

CB_NAMESPACE_END

#endif //CELESTIALBEING_UNIQUE_OBJECT_HPP
